-- 基于令牌桶算法的Redis限流脚本
-- KEYS[1] - 令牌桶key
-- KEYS[2] - 时间戳key
-- ARGV[1] - 补充速率（每秒补充的令牌数）
-- ARGV[2] - 桶容量（最大令牌数）
-- ARGV[3] - 当前时间戳
-- ARGV[4] - 请求的令牌数

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

-- 获取当前令牌数和上次补充时间
local bucket = redis.call('HMGET', tokens_key, 'tokens', 'last_refill')
local tokens = tonumber(bucket[1])
local last_refill = tonumber(bucket[2])

-- 如果桶不存在，初始化
if tokens == nil then
  tokens = capacity
  last_refill = now
end

-- 计算需要补充的令牌数
local time_passed = now - last_refill
local new_tokens = math.min(capacity, tokens + (time_passed * rate))

-- 检查是否有足够的令牌
local allowed = new_tokens >= requested
local remaining = new_tokens

if allowed then
  remaining = new_tokens - requested
end

-- 更新令牌桶状态
redis.call('HMSET', tokens_key, 'tokens', remaining, 'last_refill', now)
redis.call('EXPIRE', tokens_key, 3600) -- 1小时过期

-- 返回结果：[是否允许(1/0), 剩余令牌数]
if allowed then
  return {1, remaining}
else
  return {0, remaining}
end